import base64
import json
import os
from typing import Any

import pandas as pd
import pymysql
import httpx
import requests
from agents import OpenAIChatCompletionsModel, Agent, Runner, set_default_openai_client, function_tool, \
    set_tracing_disabled, ModelSettings
from dotenv import load_dotenv
from openai import AsyncOpenAI
from lxml import etree
import  tiktoken

load_dotenv(override=True)
##################################################################
def getAgent(instructions, tools, agent_name="deepseek_agent", handoff=None):
    if handoff is None:
        handoff = []
    client = AsyncOpenAI(api_key=os.getenv("LLM_API_KEY"), base_url=os.getenv("LLM_BASE_URL"))
    set_default_openai_client(client)
    set_tracing_disabled(True)

    deepseek_model = OpenAIChatCompletionsModel(
        openai_client=client,
        model=os.getenv("LLM_MODEL")
    )

    agent = Agent(
        name=agent_name,
        instructions=instructions,
        model=deepseek_model,
        tools=tools,
        handoffs=handoff)

    return agent


def getAgent(instructions, tools=None, agent_name="deepseek_agent", handoff=None, model_settings:ModelSettings=ModelSettings(),provider=""):
    if handoff is None:
        handoff = []
    if tools is None:
        tools = []
    client = AsyncOpenAI(api_key=os.getenv((provider + "_" if provider != "" else "") + "LLM_API_KEY"), base_url=os.getenv((provider + "_" if provider!="" else "") + "LLM_BASE_URL"))
    set_default_openai_client(client)
    set_tracing_disabled(True)

    deepseek_model = OpenAIChatCompletionsModel(
        openai_client=client,
        model=os.getenv((provider + "_" if provider!="" else "") + "LLM_MODEL")
    )

    agent = Agent(
        name=agent_name,
        instructions=instructions,
        model=deepseek_model,
        tools=tools,
        handoffs=handoff,
        model_settings=model_settings)

    return agent

async def chat(Agent):
    messages = []

    while True:
        user_input = input("请输入您的问题: ")
        if user_input == "退出":
            break
        elif user_input == '清空对话历史':
            messages = []
            print("模型回答: 清空完毕")
            continue

        messages.append({"role": "user", "content": user_input})
        result = await Runner.run(Agent, messages,max_turns=30)
        print("模型回答: ", result.final_output)
        messages = result.to_input_list()

#############################################################

async def fetch_weather(city: str) -> dict[str, Any] | None:
    """
    从 OpenWeather API 获取天气信息｡
    :param city: 城市名称(需使用英文,如 Beijing)
    :return: 天气数据字典;若出错返回包含 error 信息的字典
    """
    params = {
        "q": city,
        "appid": os.getenv("API_KEY"),
        "units": "metric",
        "lang": "zh_cn"
    }
    headers = {"User-Agent": os.getenv("USER_AGENT")}
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(os.getenv("OPENWEATHER_API_BASE"), params=params,
            headers=headers, timeout=60.0)
            response.raise_for_status()
            return response.json() # 返回字典类型
        except httpx.HTTPStatusError as e:
            return {"error": f"HTTP 错误: {e.response.status_code}"}
        except Exception as e:
            return {"error": f"请求失败: {str(e)}"}

def format_weather(data: dict[str, Any] | str) -> str:
    """
    将天气数据格式化为易读文本｡
    :param
    data: 天气数据(可以是字典或JSON字符串)
    :return: 格式化后的天气信息字符串
    """
    # 如果传入的是字符串,则先转换为字典
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except Exception as e:
            return f"无法解析天气数据: {e}"
    # 如果数据中包含错误信息,直接返回错误提示
    if "error" in data:
        return f"⚠️{data['error']}"

    # 提取数据时做容错处理
    city = data.get("name", "未知")
    country = data.get("sys", {}).get("country", "未知")
    temp = data.get("main", {}).get("temp", "N/A")
    humidity = data.get("main", {}).get("humidity", "N/A")
    wind_speed = data.get("wind", {}).get("speed", "N/A")
    # weather 可能为空列表,因此用 [0] 前先提供默认字典
    weather_list = data.get("weather", [{}])
    description = weather_list[0].get("description", "未知")
    return (
        f" {city}, {country}\n"
        f"🌡 温度: {temp}°C\n"
        f"💧 湿度: {humidity}%\n"
        f"🌬 风速: {wind_speed} m/s\n"
        f"🌤 天气: {description}\n"
    )

@function_tool
async def query_weather(city: str) -> str:
    """
    输入指定城市的英文名称, 返回今日天气查询结果｡
    :param
    city: 城市名称(需使用英文)
    :return: 格式化后的天气信息
    """
    data = await fetch_weather(city)
    result = format_weather(data)
    return result

#############################################################

@function_tool
async def writeFile(filePath: str, content: str):
    """
    写入文件
    :param filePath: 文件路径
    :param content:  文件内容
    :return:
    """

    raise RuntimeError("由于格式问题，无法写入本地文档，请调整格式后重新尝试。")


##############################################################

@function_tool
def python_inter(py_code):
    """
    专门用于执行python代码，并获取最终查询或处理结果。
    :param py_code: 字符串形式的Python代码，
    :return：代码运行的最终结果
    """
    print("正在调用python_inter工具运行Python代码...")

    g = globals()
    try:
        # 尝试如果是表达式，则返回表达式运行结果
        result = eval(py_code, g)
        return json.dumps(result, ensure_ascii=False)
        # 若报错，则先测试是否是对相同变量重复赋值
    except Exception as e:
        global_vars_before = set(g.keys())
        try:
            exec(py_code, g)
        except Exception as e:
            return f"代码执行时报错{e}"
        global_vars_after = set(g.keys())
        new_vars = global_vars_after - global_vars_before
        # 若存在新变量
        if new_vars:
            result = {var: g[var] for var in new_vars}
            print("代码已顺利执行，正在进行结果梳理...")
            return str(result)
        else:
            print("代码已顺利执行，正在进行结果梳理...")
            return "已经顺利执行代码"


######################################################################################
@function_tool
def fig_inter(py_code, fname):
    print("正在调用fig_inter工具运行Python代码...")
    import matplotlib
    import os
    import matplotlib.pyplot as plt
    import seaborn as sns
    import pandas as pd
    from IPython.display import display, Image

    g = globals()

    # 切换为无交互式后端
    current_backend = matplotlib.get_backend()
    matplotlib.use('Agg')

    # 用于执行代码的本地变量
    local_vars = {"plt": plt, "pd": pd, "sns": sns}

    # 相对路径保存目录
    pics_dir = 'pics'
    if not os.path.exists(pics_dir):
        os.makedirs(pics_dir)

    try:
        # 执行用户代码
        exec(py_code, g, local_vars)
        g.update(local_vars)

        # 获取图像对象
        fig = local_vars.get(fname, None)
        if fig:
            rel_path = os.path.join(pics_dir, f"{fname}.png")
            fig.savefig(rel_path, bbox_inches='tight')
            display(Image(filename=rel_path))
            print("代码已顺利执行，正在进行结果梳理...")
            return f"✅ 图片已保存，相对路径: {rel_path}"
        else:
            return "⚠️ 代码执行成功，但未找到图像对象，请确保有 `fig = ...`。"
    except Exception as e:
        return f"❌ 执行失败：{e}"
    finally:
        # 恢复原有绘图后端
        matplotlib.use(current_backend)

######################################################################################

@function_tool
def sql_inter(sql_query):
    """
    用于执行一段SQL代码，并最终获取SQL代码执行结果，\
    核心功能是将输入的SQL代码传输至MySQL环境中进行运行，\
    并最终返回SQL代码运行结果。需要注意的是，本函数是借助pymysql来连接MySQL数据库。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对MySQL数据库中各张表进行查询，并获得各表中的各类相关信息
    :return：sql_query在MySQL中的运行结果。
    """
    print("正在调用sql_inter工具运行SQL代码...")

    host = os.getenv('HOST')
    user = os.getenv('USER')
    mysql_pw = os.getenv('MYSQL_PW')
    db = os.getenv('DB_NAME')
    port = os.getenv('PORT')

    connection = pymysql.connect(
        host=host,
        user=user,
        passwd=mysql_pw,
        db=db,
        port=int(port),
        charset='utf8',
    )

    try:
        with connection.cursor() as cursor:
            sql = sql_query
            cursor.execute(sql)
            results = cursor.fetchall()
            print("SQL代码已顺利运行，正在整理答案...")

    finally:
        connection.close()

    return json.dumps(results)

######################################################################################

@function_tool
def extract_data(sql_query, df_name):
    """
    借助pymysql将MySQL中的某张表读取并保存到本地Python环境中。
    :param sql_query: 字符串形式的SQL查询语句，用于提取MySQL中的某张表。
    :param df_name: 将MySQL数据库中提取的表格进行本地保存时的变量名，以字符串形式表示。
    :return：表格读取和保存结果
    """
    print("正在调用extract_data工具运行SQL代码...")

    g = globals()

    host = os.getenv('HOST')
    user = os.getenv('USER')
    mysql_pw = os.getenv('MYSQL_PW')
    db = os.getenv('DB_NAME')
    port = os.getenv('PORT')

    connection = pymysql.connect(
        host=host,
        user=user,
        passwd=mysql_pw,
        db=db,
        port=int(port),
        charset='utf8',
    )

    g[df_name] = pd.read_sql(sql_query, connection)
    print("代码已顺利执行，正在进行结果梳理...")
    return "已成功创建pandas对象：%s，该变量保存了同名表格信息" % df_name


#################################################################################
def google_search(query, num_results=10, site_url=None):
    api_key = os.getenv("GOOGLE_SEARCH_API_KEY")
    cse_id = os.getenv("CSE_ID")

    url = "https://www.googleapis.com/customsearch/v1"

    # API 请求参数
    if site_url == None:
        params = {
            'q': query,
            'key': api_key,
            'cx': cse_id,
            'num': num_results
        }
    else:
        params = {
            'q': query,
            'key': api_key,
            'cx': cse_id,
            'num': num_results,
            'siteSearch': site_url
        }

    # 发送请求
    response = requests.get(url, params=params)
    response.raise_for_status()

    # 解析响应
    search_results = response.json().get('items', [])

    # 提取所需信息
    results = [{
        'title': item['title'],
        'link': item['link'],
        'snippet': item['snippet']
    } for item in search_results]

    return results


def windows_compatible_name(s, max_length=255):
    """
    将字符串转化为符合Windows文件/文件夹命名规范的名称。

    参数:
    - s (str): 输入的字符串。
    - max_length (int): 输出字符串的最大长度，默认为255。

    返回:
    - str: 一个可以安全用作Windows文件/文件夹名称的字符串。
    """

    # Windows文件/文件夹名称中不允许的字符列表
    forbidden_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*']

    # 使用下划线替换不允许的字符
    for char in forbidden_chars:
        s = s.replace(char, '_')

    # 删除尾部的空格或点
    s = s.rstrip(' .')

    # 检查是否存在以下不允许被用于文档名称的关键词，如果有的话则替换为下划线
    reserved_names = ["CON", "PRN", "AUX", "NUL", "COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8",
                      "COM9",
                      "LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9"]
    if s.upper() in reserved_names:
        s += '_'

    # 如果字符串过长，进行截断
    if len(s) > max_length:
        s = s[:max_length]

    return s


def get_search_text(q, url):
    cookie = os.getenv('search_cookie')
    user_agent = os.getenv('search_user_agent')

    code_ = False
    headers = {
        'authority': 'www.zhihu.com',
        'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
        'cache-control': 'max-age=0',
        'cookie': cookie,
        'upgrade-insecure-requests': '1',
        'user-agent': user_agent,
    }

    # 普通问答地址
    if 'zhihu.com/question' in url:
        res = requests.get(url, headers=headers).text
        res_xpath = etree.HTML(res)
        title = res_xpath.xpath('//div/div[1]/div/h1/text()')[0]
        text_d = res_xpath.xpath('//div/div/div/div[2]/div/div[2]/div/div/div[2]/span[1]/div/div/span/p/text()')

    # 专栏地址
    elif 'zhuanlan' in url:
        headers['authority'] = 'zhaunlan.zhihu.com'
        res = requests.get(url, headers=headers).text
        res_xpath = etree.HTML(res)

        titles = res_xpath.xpath('//*[@id="root"]/div/main/div/div[2]/div[1]/div/article/header/h1/text()')
        if len(titles)>0:
            title = titles[0]
        else:
            title = ""
        #text_d = res_xpath.xpath('//*[@id="root"]/div/main/div/div[2]/div[1]/div/article/div[1]')[0].text()
        text_d_array = res_xpath.xpath('//*[@id="root"]/div/main/div/div[2]/div[1]/div/article/div[1]/div/div/div[2]//text()')
        text_d = ''.join(text_d_array)
        code_ = res_xpath.xpath('//*[@id="root"]/div/main/div/div[2]/div[1]/div/article//pre/code/text()')
        # 特定回答的问答网址
    elif 'answer' in url:
        res = requests.get(url, headers=headers).text
        res_xpath = etree.HTML(res)
        title = res_xpath.xpath('//div/div[1]/div/h1/text()')[0]
        text_d = res_xpath.xpath('//div[1]/div/div[3]/div/div/div/div[2]/span[1]/div/div/span/p/text()')

    if title == None:
        return None

    else:
        title = windows_compatible_name(title)

        # 创建问题答案正文
        text = ''
        for t in text_d:
            txt = str(t).replace('\n', ' ')
            text += txt

        # 如果有code，则将code追加到正文的追后面
        if code_:
            for c in code_:
                co = str(c).replace('\n', ' ')
                text += co

        encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
        json_data = [
            {
                "link": url,
                "title": title,
                "content": text,
                "tokens": len(encoding.encode(text))
            }
        ]

        # 自动创建目录，如果不存在的话
        dir_path = f'./auto_search/{q}'
        os.makedirs(dir_path, exist_ok=True)

        with open('./auto_search/%s/%s.json' % (q, title), 'w') as f:
            json.dump(json_data, f)

        return title


def get_search_result(q):
    """
    当你无法回答某个问题时，调用该函数，能够获得答案
    :param q: 必选参数，询问的问题，字符串类型对象
    :return：某问题的答案，以字符串形式呈现
    """
    # 默认搜索返回5个答案
    results = google_search(query=q, num_results=5, site_url='https://zhihu.com/')

    # 创建对应问题的子文件夹
    folder_path = './auto_search/%s' % q
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)

    # 单独提取links放在一个list中
    num_tokens = 0
    content = ''
    for item in results:
        url = item['link']
        title = get_search_text(q, url)
        with open('./auto_search/%s/%s.json' % (q, title), 'r') as f:
            jd = json.load(f)
        num_tokens += jd[0]['tokens']
        if num_tokens <= 12000:
            # print(jd[0]['content'])
            content += jd[0]['content']
        else:
            break
    return (content)

#########################################################################################
def save_markdown_to_file(content: str, filename_hint: str, directory="research_task"):
    # 在当前项目目录下创建 research_task 文件夹
    save_dir = os.path.join(os.getcwd(), directory)

    # 如果目录不存在则创建
    os.makedirs(save_dir, exist_ok=True)

    # 创建文件名（取前8个字符并加上...）
    filename = f"{filename_hint[:8]}....md"

    # 完整文件路径
    file_path = os.path.join(save_dir, filename)

    # 将内容保存为Markdown文档
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(content)

    print(f"文件已成功保存到：{file_path}")


#########################################################################################
def get_github_readme(dic):
    github_token = os.getenv('GITHUB_TOKEN')
    user_agent = os.getenv('search_user_agent')

    owner = dic['owner']
    repo = dic['repo']

    headers = {
        "Authorization": github_token,
        "User-Agent": user_agent
    }

    response = requests.get(f"https://api.github.com/repos/{owner}/{repo}/readme", headers=headers)

    readme_data = response.json()
    encoded_content = readme_data.get('content', '')
    decoded_content = base64.b64decode(encoded_content).decode('utf-8')

    return decoded_content

def extract_github_repos(search_results):
    # 使用列表推导式筛选出项目主页链接
    repo_links = [result['link'] for result in search_results if '/issues/' not in result['link'] and '/blob/' not in result['link'] and 'github.com' in result['link'] and len(result['link'].split('/')) == 5]

    # 从筛选后的链接中提取owner和repo
    repos_info = [{'owner': link.split('/')[3], 'repo': link.split('/')[4]} for link in repo_links]

    return repos_info


def get_search_text_github(q, dic):
    title = dic['owner'] + '_' + dic['repo']
    title = windows_compatible_name(title)

    # 创建问题答案正文
    text = get_github_readme(dic)

    # 写入本地json文件
    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
    json_data = [
        {
            "title": title,
            "content": text,
            "tokens": len(encoding.encode(text))
        }
    ]

    # 自动创建目录，如果不存在的话
    dir_path = f'./auto_search/{q}'
    os.makedirs(dir_path, exist_ok=True)

    with open('./auto_search/%s/%s.json' % (q, title), 'w') as f:
        json.dump(json_data, f)

    return title

@function_tool
def get_answer(q):
    """
    当你无法回答某个问题时，调用该函数，能够获得答案
    :param q: 必选参数，询问的问题，字符串类型对象
    :param g: g，字符串形式变量，表示环境变量，无需设置，保持默认参数即可
    :return：某问题的答案，以字符串形式呈现
    """
    # 默认搜索返回5个答案
    print('正在接入谷歌搜索，查找和问题相关的答案...')
    results = google_search(query=q, num_results=5, site_url='https://zhihu.com/')

    # 创建对应问题的子文件夹
    folder_path = './auto_search/%s' % q
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)

    # 单独提取links放在一个list中
    num_tokens = 0
    content = ''
    for item in results:
        url = item['link']
        print('正在检索：%s' % url)
        title = get_search_text(q, url)
        with open('./auto_search/%s/%s.json' % (q, title), 'r') as f:
            jd = json.load(f)
        num_tokens += jd[0]['tokens']
        if num_tokens <= 12000:
            # print(jd[0]['content'])
            content += jd[0]['content']
        else:
            break
    return (content)


#######################
@function_tool
def get_answer_github(q):
    """
    当你无法回答某个问题时，调用该函数，能够获得答案
    :param q: 必选参数，询问的问题，字符串类型对象
    :param g: g，字符串形式变量，表示环境变量，无需设置，保持默认参数即可
    :return：某问题的答案，以字符串形式呈现
    """
    # 调用转化函数，将用户的问题转化为更适合在GitHub上搜索的关键词
    # q = convert_keyword_github(q)

    # 默认搜索返回5个答案
    print('正在接入谷歌搜索，查找和问题相关的答案...')
    search_results = google_search(query=q, num_results=5, site_url='https://github.com/')
    results = extract_github_repos(search_results)

    # 创建对应问题的子文件夹
    folder_path = './auto_search/%s' % q
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)

    print('正在读取相关项目说明文档...')
    num_tokens = 0
    content = ''

    for dic in results:
        title = get_search_text_github(q, dic)
        with open('./auto_search/%s/%s.json' % (q, title), 'r') as f:
            jd = json.load(f)
        num_tokens += jd[0]['tokens']
        if num_tokens <= 12000:
            content += jd[0]['content']
        else:
            break
    print('正在进行最后的整理...')
    return (content)

#######################################################################
def print_code_if_exists(function_args):
    """
    如果存在代码片段，则打印代码
    """

    def convert_to_markdown(code, language):
        return f"```{language}\n{code}\n```"

    # 如果是SQL，则按照Markdown中SQL格式打印代码
    if function_args.get('sql_query'):
        code = function_args['sql_query']
        markdown_code = convert_to_markdown(code, 'sql')
        print("即将执行以下代码：")
        print(markdown_code)

    # 如果是Python，则按照Markdown中Python格式打印代码
    elif function_args.get('py_code'):
        code = function_args['py_code']
        markdown_code = convert_to_markdown(code, 'python')
        print("即将执行以下代码：")
        #display(Markdown(markdown_code))
        print(markdown_code)